import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.cluster import KMeans
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib
matplotlib.use('Agg')

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False

data = pd.read_csv(r"data/electricity_info.csv")

monthly_usage = data.groupby(['user_index', 'power_month'])['basic_power'].sum().unstack(fill_value=0)
monthly_usage.columns = [f"month_{col}" for col in monthly_usage.columns]
monthly_usage.to_csv(r'data/monthly_usage.csv')
static_features = data.drop_duplicates(subset='user_index')[[
    'user_index', 'user_age_range_id', 'user_gender', 'user_city_level_id',
    'user_wage_range_id', 'user_professional_status'
]]
static_features.to_csv(r'data/static_features.csv')
merged = pd.merge(static_features, monthly_usage, on='user_index')
merged.to_csv(r'data/merged.csv')
ohe_cols = ['user_gender', 'user_professional_status']
num_cols = ['user_age_range_id', 'user_city_level_id', 'user_wage_range_id'] + list(monthly_usage.columns)

preprocessor = ColumnTransformer([
    ('num', StandardScaler(), num_cols),
    ('cat', OneHotEncoder(sparse_output=False), ohe_cols)
])

pipeline = Pipeline([
    ('preprocess', preprocessor),
    ('cluster', KMeans(n_clusters=4, random_state=42))
])

pipeline.fit(merged.drop(columns=['user_index']))

merged['cluster'] = pipeline.named_steps['cluster'].labels_

X_reduced = PCA(n_components=2).fit_transform(
    pipeline.named_steps['preprocess'].transform(merged.drop(columns=['user_index', 'cluster']))
)

merged['pca_1'] = X_reduced[:, 0]
merged['pca_2'] = X_reduced[:, 1]

print(merged[['user_index', 'cluster'] + list(monthly_usage.columns)].head())
merged.to_csv(r'data/last_merged.csv')

plt.figure(figsize=(8, 6))
sns.scatterplot(data=merged, x='pca_1', y='pca_2', hue='cluster', palette='Set2')
plt.title('模拟燃气使用习惯用户聚类图')
plt.xlabel('PCA 1 用气总量 + 城市等级 + 收入')
plt.ylabel('PCA 2 用气稳定性 + 职业状态')
plt.legend(title='Cluster')
plt.grid(True)
plt.tight_layout()
plt.savefig('cluster_result.png', dpi=150)

usage_summary = merged.groupby('cluster')[monthly_usage.columns].mean()
print(usage_summary)
usage_summary.to_csv(r'data/usage_summary.csv')
profile_summary = merged.groupby('cluster')[[
    'user_age_range_id', 'user_city_level_id', 'user_wage_range_id'
]].agg(['mean', 'median'])

plt.figure(figsize=(8, 6))
sns.scatterplot(data=merged, x='pca_1', y='pca_2', hue='cluster', palette='Set2')
plt.title('模拟燃气使用习惯用户聚类图')
plt.xlabel('PCA 1 用气总量 + 城市等级 + 收入')
plt.ylabel('PCA 2 用气稳定性 + 职业状态')
plt.legend(title='Cluster')
plt.grid(True)
plt.tight_layout()
plt.savefig('figures/cluster_result.png', dpi=150)
plt.clf()

plt.figure(figsize=(10, 6))
usage_summary.T.plot()
plt.title("各类用户月度用气平均曲线")
plt.xlabel("月份")
plt.ylabel("用气量")
plt.legend(title="Cluster")
plt.grid(True)
plt.tight_layout()
plt.savefig("figures/monthly_usage_by_cluster.png", dpi=150)
plt.clf()

plt.figure(figsize=(8, 6))
sns.countplot(data=merged, x='user_age_range_id', hue='cluster', palette='Set2')
plt.title("不同年龄段用户在各聚类中的分布")
plt.xlabel("年龄段编号")
plt.ylabel("用户数量")
plt.legend(title="Cluster")
plt.tight_layout()
plt.savefig("figures/age_distribution_by_cluster.png", dpi=150)
plt.clf()




